import cv2 as cv
import numpy as np
import apriltag
from math import *
from uarm.wrapper import SwiftAPI

swift = SwiftAPI()
swift.reset()

# 相机内参矩阵
mtx = np.load('./data/camera_param.npz')['mtx']
dist = np.load('./data/camera_param.npz')['dist']

# 定义一个检测器（使用字典“tag36h10”）
at_detector = apriltag.Detector(apriltag.DetectorOptions(families='tag36h11 tag36h10'))

def eulerAngles2RotationMatrix(theta):
    """
    欧拉角转换旋转矩阵
    :param theta: 欧拉角角度
    :return: 旋转矩阵
    """
    R_x = np.array([[1, 0, 0],
                    [0, cos(theta[0]), sin(theta[0])],
                    [0, -sin(theta[0]), cos(theta[0])]
                    ])
    R_y = np.array([[cos(theta[1]), 0, sin(theta[1])],
                    [0, 1, 0],
                    [-sin(theta[1]), 0, cos(theta[1])]
                    ])
    R_z = np.array([[cos(theta[2]), -sin(theta[2]),0],
                    [sin(theta[2]), cos(theta[2]),0],
                    [0, 0, 1]
                    ])
    R = R_z@R_y@R_x
    return R

def rotationMatrixToEulerAngles(R):
    """
    旋转矩阵转换欧拉角
    :param R: 旋转矩阵
    :return: 欧拉角
    """
    sy = sqrt(R[0, 0] * R[0, 0] + R[1, 0] * R[1, 0])
    singular = sy < 1e-6
    if not singular:
        x = atan2(R[2, 1], R[2, 2])
        y = atan2(-R[2, 0], sy)
        z = atan2(R[1, 0], R[0, 0])
    else:
        x = atan2(-R[1, 2], R[1, 1])
        y = atan2(-R[2, 0], sy)
        z = 0
    return np.array([x, y, z])


def get_tag2end(point):
    """
    获取二维码到手臂末端
    :param point:
    :return:
    """

    # 计算绕z轴旋转角
    obj= point[:,:2]
    delta = (obj[2]-obj[3])
    r = sqrt(delta@delta.T)
    angle = np.arcsin(delta[0]/r)
    # 计算变换后的坐标
    position = swift.get_position()
    R=eulerAngles2RotationMatrix([0,0,angle])
    t = R @ np.array(position) - R @ point[3]
    return R,t

def get_img2tag(image,tag_size):
    img2tag_p = []
    cam_params = [mtx[0, 0], mtx[1, 1], mtx[0, 2], mtx[1, 2]]
    gray = cv.cvtColor(image,cv.COLOR_BGR2GRAY)
    tags = at_detector.detect(gray)
    for tag in tags:
        M, e1, e2 = at_detector.detection_pose(tag, cam_params, tag_size)
        img2tag_p.append([tag.tag_id,tag.center,M])
    return img2tag_p

def get_base2end():
    angle = swift.get_polar()[1]
    position = swift.get_position()
    R = eulerAngles2RotationMatrix([0,0,(angle-90)*pi/180])
    t = position
    return R,t

def get_base2tag_p(tags,id,R,t):
    ret = False
    dst = []
    for tag in tags:
        if tag[0]==id:
            pix2img_t = np.linalg.inv(mtx) @ np.insert(tag[1], 2, 1)
            img2tag_t = (tag[2] @ np.insert(pix2img_t, 3, 1))[:3]
            base2end_R, base2end_t = get_base2end()
            dst = base2end_t + base2end_R @ R @ img2tag_t + base2end_R @ -R@t
            ret = True
            print(dst)
    return ret,dst

def grad_id(image,id):
    with np.load('./data/eyehand_Matrix.npz') as X:
        R, t = [X[i] for i in ('R', 't')]
    tags = get_img2tag(image, 20)
    ret, dst_p = get_base2tag_p(tags, id, R, t)
    if ret == True:
        swift.set_position(dst_p[0], dst_p[1], int(dst_p[2]+10))
        swift.set_position(dst_p[0], dst_p[1], int(dst_p[2]))
        swift.set_position(dst_p[0], dst_p[1], int(dst_p[2]))
    else:
        print("没有识别到id%d",id)


def waitMouse(event,x,y,flags,param):
    if event == cv.EVENT_LBUTTONDOWN:
        print(x,y)


if __name__ == "__main__":
    cap = cv.VideoCapture(4)
    cv.namedWindow("image")
    cv.setMouseCallback("image",waitMouse)
    while cap.isOpened():
        _,frame = cap.read()
        cv.imshow("image",frame)
        key = cv.waitKey(10)
        if key == 27:
            break
    cap.release()
    cv.destroyAllWindows()

